package com.auditlog.datasource.context;

import cn.hutool.core.util.IdUtil;
import com.auditlog.datasource.context.ConnectionContext;
import com.auditlog.datasource.context.ExecutionContext;
import lombok.extern.slf4j.Slf4j;

import java.util.HashMap;
import java.util.Map;

/**
 * @author Zhiyang.Zhang
 * @version 1.0
 * @date 2022/11/28 20:32
 */
@Slf4j
public class ContextHolder {

    public static final String EXECUTE_CONTEXT = "executeContext";

    public static final String CONNECTION_CONTEXT = "connectionContext";

    private static final ThreadLocal<Map<String, Object>> CONTEXT_HOLDER = ThreadLocal.withInitial(() -> new HashMap<>());


    public static void bind(String key, Object context) {
        CONTEXT_HOLDER.get().put(key, context);
    }

    public static Object getObject(String key) {
        return CONTEXT_HOLDER.get().get(key);
    }

    public static ExecutionContext getExecuteContext() {
        ExecutionContext executionContext = (ExecutionContext) CONTEXT_HOLDER.get().get(EXECUTE_CONTEXT);
        if (executionContext == null) {
            executionContext = new ExecutionContext();
            bindExecutionContext(executionContext);
        }
        return executionContext;
    }

    public static ConnectionContext getConnectionContext() {
        ConnectionContext connectionContext = (ConnectionContext) CONTEXT_HOLDER.get().get(CONNECTION_CONTEXT);
        if (connectionContext == null) {
            connectionContext = new ConnectionContext();
            bindConnectionContext(connectionContext);
            getExecuteContext().setOperationId(IdUtil.getSnowflakeNextIdStr());
        }
        return connectionContext;
    }

    public static void bindExecutionContext(ExecutionContext executionContext) {
        CONTEXT_HOLDER.get().put(EXECUTE_CONTEXT, executionContext);
    }

    public static void bindConnectionContext(ConnectionContext connectionContext) {
        CONTEXT_HOLDER.get().put(CONNECTION_CONTEXT, connectionContext);
    }

    public static void unbindExecutionContext() {
        CONTEXT_HOLDER.get().remove(EXECUTE_CONTEXT);
    }

    public static void unbindConnectionContext() {
        CONTEXT_HOLDER.get().remove(CONNECTION_CONTEXT);
    }

    public static void clear() {
        CONTEXT_HOLDER.remove();
    }

    public static void reset() {
        getConnectionContext().resetAll();
        getExecuteContext().reset();
    }
}
